import evaluate
import numpy as np
import torch
from datasets import DatasetDict, load_dataset
from transformers import AutoTokenizer, AutoModelForTokenClassification, TrainingArguments, Trainer,\
	DataCollatorForTokenClassification, pipeline
torch.cuda.empty_cache()
torch.cuda.set_device(0)

### 中文命名实体识别数据集（人民日报语料）
## 加载数据集
ner_datasets = load_dataset("lansinuote/peoples-daily-ner", cache_dir = "./ner_data")
print(ner_datasets)
# 获取训练集单个样本
print(ner_datasets["train"][0])
# 训练集的特征结构定义
print(ner_datasets["train"].features)
# 从数据集特征定义中提取命名实体标签的名称列表
# 就是找到列表 ['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC']
label_list = ner_datasets["train"].features["ner_tags"].feature.names
print(label_list)
print("------------------------------------1--------------------------------------")
# 数据集预处理
#加载预训练的中文 MacBERT 模型的分词器
tokenizer = AutoTokenizer.from_pretrained("hfl/chinese-macbert-base")
# 对训练集第一个样本的 tokens 进行分词处理,这会将中文字符转换为模型可以理解的 token ID
# 对于已经做好tokenize的数据，要指定is_split_into_words参数为True
tokenizer(ner_datasets["train"][0]["tokens"], is_split_into_words = True)
# 对英文文本 "interesting word" 进行分词处理
res = tokenizer("interesting word")
print(res)
print(res.word_ids())

#word_ids 映射关系
def process_function(examples):
	# 使用预训练的tokenizer对输入的tokens进行分词处理
	# 参数说明：
	# - examples["tokens"]: 输入的文本tokens列表
	# - max_length=64: 设置序列最大长度为64，超过的部分会被截断
	# - truncation=True: 允许截断超过最大长度的序列
	# - is_split_into_words=True: 指示输入已经是分词好的单词列表格式
	# 当 ner_datasets.map处理训练集时，会自动构造 examples 参数:
	tokenized_exmaples = tokenizer(examples["tokens"], max_length = 32, truncation = True, is_split_into_words = True)
	# 初始化标签列表，用于存储处理后的标签
	labels = []
	# 遍历每一个样本的ner_tags标签
	# i: 样本索引
	# label: 当前样本的标签列表
	for i, label in enumerate(examples["ner_tags"]):
		# 获取当前样本分词后的word_ids映射关系
		# word_ids表示每个token对应原始单词的索引，None表示特殊token(如[CLS],[SEP]
		word_ids = tokenized_exmaples.word_ids(batch_index = i)
		# 初始化当前样本的标签ID列表
		label_ids = []
		# 遍历每个word_id
		for word_id in word_ids:
			if word_id is None:
				# 如果word_id为None，说明是特殊token(如[CLS],[SEP])，标记为-100(在计算损失时会被忽略)
				label_ids.append(-100)
			else:
				# 否则使用原始标签，根据word_id找到对应的标签值
				# 这样处理是因为一个单词可能被分成多个token，它们共享同一个标签
				label_ids.append(label[word_id])
		# 将处理好的标签列表添加到总标签列表中
		labels.append(label_ids)
	# 将处理后的标签添加到分词结果中，作为模型训练的标签
	tokenized_exmaples["labels"] = labels
	return tokenized_exmaples


# 对数据集应用预处理函数，batched=True表示以批次方式处理数据
# process_function是自定义的数据预处理函数
# 这会返回一个新的DatasetDict对象，包含处理后的训练集、验证集和测试集
tokenized_datasets = ner_datasets.map(process_function, batched = True)
print(tokenized_datasets)
print(tokenized_datasets["train"][0])
print("------------------------------------2--------------------------------------")
# 创建模型
# 对于所有的非二分类任务，切记要指定num_labels，否则就会device错误
model = AutoModelForTokenClassification.from_pretrained("hfl/chinese-macbert-base", num_labels = len(label_list))
print(model.config)


# 确保模型的所有权重都是连续的
def ensure_weights_contiguous(model):
	for name, param in model.named_parameters():
		if not param.is_contiguous():
			print(f"Making {name} contiguous.")
			param.data = param.data.contiguous()


ensure_weights_contiguous(model)
print(model.config.num_labels)
# print("------------------------------------3--------------------------------------")
# 创建评估函数
# 本地的加载方式，无需额外下载
seqeval = evaluate.load("seqeval_metric.py")
print(seqeval)


def eval_metric(pred):
    # 解包预测结果，pred是一个元组，包含模型的预测输出和真实标签
    predictions, labels = pred
    # 对预测结果取最大值的索引，得到每个位置最可能的标签ID
    predictions = np.argmax(predictions, axis=-1)

    # 将预测的标签ID转换为原始的字符串标签
    # 使用列表推导式遍历每个样本的预测结果和真实标签
    # 只保留标签不为-100的位置（-100是特殊标记，在计算损失时被忽略）
    true_predictions = [
        [label_list[p] for p, l in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]

    # 同样地，将真实的标签ID转换为原始的字符串标签
    # 只保留标签不为-100的位置
    true_labels = [
        [label_list[l] for p, l in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]

    # 使用seqeval库计算评估指标
    # predictions: 模型预测的标签序列
    # references: 真实的标签序列
    # mode="strict": 严格模式，要求实体边界完全匹配才算正确
    # scheme="IOB2": 使用IOB2标注方案（B表示实体开始，I表示实体内部，O表示其他）
    result = seqeval.compute(predictions=true_predictions, references=true_labels, mode="strict", scheme="IOB2")

    # 返回F1分数作为评估指标
    # result["overall_f1"]是seqeval计算出的整体F1分数
    return {
        "f1": result["overall_f1"]
    }

print("------------------------------------4--------------------------------------")
# 配置训练参数
args = TrainingArguments(
	output_dir = "models_for_ner",  # 模型输出目录
	per_device_train_batch_size = 64,  # 每个设备上的训练批量大小
	per_device_eval_batch_size = 128,  # 每个设备上的评估批量大小
	gradient_accumulation_steps = 4,  # 梯度累积步数
	gradient_checkpointing = True,  # 启用梯度检查点以节省内存
	gradient_checkpointing_kwargs = {"use_reentrant": False},
	fp16 = True,  # 启用混合精度训练以节省显存
	optim = "adafactor",  # 使用 Adafactor 优化器
	eval_strategy = "epoch",  # 评估策略：每个训练周期结束时进行一次评估
	save_strategy = "epoch",  # 保存策略：每个训练周期结束时保存一次模型
	metric_for_best_model = "f1",  # 用于确定最佳模型的指标
	load_best_model_at_end = True,  # 训练结束后加载最佳模型
	logging_steps = 50,  # 日志记录间隔：每 50 个步骤输出一次日志
	num_train_epochs = 1  # 总训练轮数
)
print(args)
print("------------------------------------5--------------------------------------")
# 创建训练器
trainer = Trainer(
	model = model,
	args = args,
	tokenizer = tokenizer,
	train_dataset = tokenized_datasets["train"],
	eval_dataset = tokenized_datasets["validation"],
	compute_metrics = eval_metric,
	data_collator = DataCollatorForTokenClassification(tokenizer = tokenizer)
)
print(trainer)
print("------------------------------------6--------------------------------------")
# 模型训练
trainer.train()
trainer.evaluate(eval_dataset = tokenized_datasets["test"])
print("------------------------------------7--------------------------------------")
# 模型预测
# 使用pipeline进行推理，要指定id2label
model.config.id2label = {idx: label for idx, label in enumerate(label_list)}
print(model.config)
# 如果模型是基于GPU训练的，那么推理时要指定device
# 对于NER任务，可以指定aggregation_strategy为simple，得到具体的实体的结果，而不是token的结果
ner_pipe = pipeline("token-classification", model = model, tokenizer = tokenizer, device = 0,
	aggregation_strategy = "simple")
res = ner_pipe("小明在北京上班")
print(res)
# 根据start和end取实际的结果
ner_result = {}
x = "浩浩在北京上班"
for r in res:
	if r["entity_group"] not in ner_result:
		ner_result[r["entity_group"]] = []
	ner_result[r["entity_group"]].append(x[r["start"]: r["end"]])

print(ner_result)
print("------------------------------------8--------------------------------------")